# coding=utf-8
import math
import random
from math import sqrt

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import irfft


def wgn(x, snr):
    snr = 10 ** (snr / 10.0)
    x_power = np.sum(x ** 2) / len(x)
    n_power = x_power / snr
    return np.random.randn(len(x)) * np.sqrt(n_power)


def noise(N, color='white', state=None):
    """Noise generator.

    :param N: Amount of samples.
    :param color: Color of noise.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    """
    try:
        return _noise_generators[color](N, state)
    except KeyError:
        raise ValueError("Incorrect color.")


def white(N, state=None):
    """
    White noise.

    :param N: Amount of samples.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    White noise has a constant power density. It's narrowband spectrum is therefore flat.
    The power in white noise will increase by a factor of two for each octave band,
    and therefore increases with 3 dB per octave.
    """
    state = np.random.RandomState() if state is None else state
    return state.randn(N)


def pink(N, state=None):
    """
    Pink noise.

    :param N: Amount of samples.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    Pink noise has equal power in bands that are proportionally wide.
    Power density decreases with 3 dB per octave.

    """
    # This method uses the filter with the following coefficients.
    # b = np.array([0.049922035, -0.095993537, 0.050612699, -0.004408786])
    # a = np.array([1, -2.494956002, 2.017265875, -0.522189400])
    # return lfilter(B, A, np.random.randn(N))
    # Another way would be using the FFT
    # x = np.random.randn(N)
    # X = rfft(x) / N
    state = np.random.RandomState() if state is None else state
    uneven = N % 2
    X = state.randn(N // 2 + 1 + uneven) + 1j * state.randn(N // 2 + 1 + uneven)
    S = np.sqrt(np.arange(len(X)) + 1.)  # +1 to avoid divide by zero
    y = (irfft(X / S)).real
    if uneven:
        y = y[:-1]
    return normalize(y)


def blue(N, state=None):
    """
    Blue noise.

    :param N: Amount of samples.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    Power increases with 6 dB per octave.
    Power density increases with 3 dB per octave.

    """
    state = np.random.RandomState() if state is None else state
    uneven = N % 2
    X = state.randn(N // 2 + 1 + uneven) + 1j * state.randn(N // 2 + 1 + uneven)
    S = np.sqrt(np.arange(len(X)))  # Filter
    y = (irfft(X * S)).real
    if uneven:
        y = y[:-1]
    return normalize(y)


def brown(N, state=None):
    """
    Violet noise.

    :param N: Amount of samples.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    Power decreases with -3 dB per octave.
    Power density decreases with 6 dB per octave.

    """
    state = np.random.RandomState() if state is None else state
    uneven = N % 2
    X = state.randn(N // 2 + 1 + uneven) + 1j * state.randn(N // 2 + 1 + uneven)
    S = (np.arange(len(X)) + 1)  # Filter
    y = (irfft(X / S)).real
    if uneven:
        y = y[:-1]
    return normalize(y)


def violet(N, state=None):
    """
    Violet noise. Power increases with 6 dB per octave.

    :param N: Amount of samples.
    :param state: State of PRNG.
    :type state: :class:`np.random.RandomState`

    Power increases with +9 dB per octave.
    Power density increases with +6 dB per octave.

    """
    state = np.random.RandomState() if state is None else state
    uneven = N % 2
    X = state.randn(N // 2 + 1 + uneven) + 1j * state.randn(N // 2 + 1 + uneven)
    S = (np.arange(len(X)))  # Filter
    y = (irfft(X * S)).real
    if uneven:
        y = y[:-1]
    return normalize(y)


_noise_generators = {
    'white': white,
    'pink': pink,
    'blue': blue,
    'brown': brown,
    'violet': violet,
}


def ms(x):
    """Mean value of signal `x` squared.

    :param x: Dynamic quantity.
    :returns: Mean squared of `x`.

    """
    return (np.abs(x) ** 2.0).mean()


def normalize(y, x=None):
    """normalize power in y to a (standard normal) white noise signal.

    Optionally normalize to power in signal `x`.

    #The mean power of a Gaussian with :math:`\\mu=0` and :math:`\\sigma=1` is 1.
    """
    # return y * np.sqrt( (np.abs(x)**2.0).mean() / (np.abs(y)**2.0).mean() )
    if x is not None:
        x = ms(x)
    else:
        x = 1.0
    return y * np.sqrt(x / ms(y))


def noise_generator(N=44100, color='white', state=None):
    """Noise generator.

    :param state:
    :param N: Amount of unique samples to generate.
    :param color: Color of noise.

    Generate `N` amount of unique samples and cycle over these samples.

    """
    # yield from itertools.cycle(noise(N, color)) # Python 3.3
    for sample in np.itertools.cycle(noise(N, color, state)):
        yield sample


def heaviside(N):
    """Heaviside.

    Returns the value 0 for `x < 0`, 1 for `x > 0`, and 1/2 for `x = 0`.
    """
    return 0.5 * (np.sign(N) + 1)


def normal_distribution(N, sigma=0.065, top=1):
    mu = N / 2.0
    sigma = sigma / (N / 100.0)

    def point(x_in):
        return 1/sqrt(2 * math.pi)*sigma*math.exp(-1 * math.pow(x_in - mu, 2) / 2 * math.pow(sigma, 2)) / 0.025 * top / (100.0 / N)

    ufunc_normal_distribute = np.frompyfunc(point, 1, 1)
    xx_out = np.arange(0, N-1, 1)
    xx_out = ufunc_normal_distribute(xx_out)
    return xx_out


if __name__ == "__main__":

    x = normal_distribution(11)
    plt.plot(x)
    plt.show()
    exit(0)
    plt.rcParams.update({'figure.autolayout': True})

    # for i in range(1):
    #     t = np.arange(0, 100) * 0.1
    #     xx = np.sin(t)
    #     n = wgn(xx, 6)
    #     f_main, (ax_main) = plt.subplots(figsize=(10, 4), nrows=1)
    #     ax_main.plot(n)
    #     plt.show()
    #     plt.clf()

    for i in range(10):
        vol = .30
        lag = 10
        df = pd.DataFrame(np.random.randn(100) * sqrt(vol) * sqrt(1 / 252.)).cumsum()
        a = pd.rolling_mean(df, lag)
        plt.plot(a.values)
        plt.show()
        plt.clf()

    for i in range(10):
        size = 52
        x_white = white(size)
        x_pink = pink(size)
        x3 = blue(size)
        x4 = brown(size)
        x_violet = violet(size)
        # function_list = [white(100), pink(100), blue(100), brown(100), violet(100)]
        # name_list = ['']

        f_main, ax_list = plt.subplots(figsize=(10 * size / 50, 8), nrows=6)  # (ax_1, ax_2, ax_3, ax_4, ax_5)

        ax_list[0].plot(x_white, label="white")
        ax_list[1].plot(x_pink, label="pink")
        ax_list[2].plot(x3, label="blue")
        ax_list[3].plot(x4, label="brown")
        ax_list[4].plot(x_violet, label="violet")

        for i in range(size / 4):
            seed = random.uniform(2, 4)

        ax_list[5].plot(x_pink + x_violet, label="x_pink + x_violet")

        for j in range(len(ax_list)):

            handles, labels = ax_list[j].get_legend_handles_labels()
            # reverse the order
            ax_list[j].legend(handles[::-1], labels[::-1])

        plt.show()
